#include <cstring>
#include <functional>

#include "tinyndarray.h"
#include "unary.h"

namespace tinyndarray {

NdArray ApplyUnaryOp(const NdArray& x, UnaryElementOp op) {
    NdArray ret(x.dtype(), x.shape());
    ApplyOpSimple(ret, x, op);
    return ret;
}

NdArray ApplyUnaryOpInplace(NdArray&& x, UnaryElementOp op) {
    ApplyOpSimple(x, x, op);
    return std::move(x);
}

template <typename T>
void neg_op(void* out, const void* a) {
    T* out_typed = static_cast<T*>(out);
    const T* a_typed = static_cast<const T*>(a);
    *out_typed = -*a_typed;
}

NdArray Neg(const NdArray& x) {
    UnaryElementOp typed_op;
    switch (x.dtype().tag) {
        case FP32: {
            typed_op = neg_op<float>;
            break;
        }
        case INT32: {
            typed_op = neg_op<int>;
            break;
        }
        case FP64: {
            typed_op = neg_op<double>;
            break;
        }
        default: throw std::runtime_error("Unsupported data type for Neg");
    }

    return ApplyUnaryOp(x, typed_op);
}

NdArray operator-(const NdArray& x) {
    return Neg(x);
}

void exp_op_float(void* out, const void* a) {
    float* out_typed = static_cast<float*>(out);
    const float* a_typed = static_cast<const float*>(a);
    *out_typed = expf(*a_typed);
}

void exp_op_double(void* out, const void* a) {
    double* out_typed = static_cast<double*>(out);
    const double* a_typed = static_cast<const double*>(a);
    *out_typed = exp(*a_typed);
}


NdArray Exp(const NdArray& x) {
    UnaryElementOp typed_op;
    switch (x.dtype().tag) {
        case FP32: {
            typed_op = exp_op_float;
            break;
        }
        case FP64: {
            typed_op = exp_op_double;
            break;
        }
        default: throw std::runtime_error("Unsupported data type for Exp");
    }

    return ApplyUnaryOp(x, typed_op);
}


}  // namespace tinyndarray